import pyqtgraph as pg
import sys
from PyQt5.QtWidgets import QWidget, QApplication, QFrame, QGridLayout, QLabel, QPushButton, QVBoxLayout
from PyQt5.QtCore import Qt, QTimer, pyqtSignal
import serial
from ctypes import CDLL
from utils.ecg_utils import EKG_QRS_detect
import numpy as np
import threading
import time


class HeartSensor(object):
    """
    A Heart Sensor Class.

    sensor = HeartSensor(dev='/dev/ttyUSB0', baud_rate=115200)
    data_list = sensor.read()
    """
    def __init__(self, dev, baud_rate):
        self.ser = serial.Serial(dev, baud_rate)

    def read(self, num=100):
        raw = list(self.ser.read(num))
        idx = 0
        for idx, item in enumerate(raw):
            if item == 170 and raw[idx + 1] == 170:
                break
        start_idx, frame_num = idx, (len(raw) - idx) // 10
        result = [raw[start_idx + i * 10 + 6: start_idx + i * 10 + 10] for i in range(frame_num)]
        result = [(item[0] << 24) | (item[1] << 16) | (item[2] << 8) | item[3] for item in result]
        return result

    def read_one(self):
        while True:
            if self.ser.read()[0] == 170:
                if self.ser.read()[0] == 170:
                    d1, d2, d3, d4, d5, d6, d7, d8 = list(self.ser.read(8))
                    result = (d5 << 24) | (d6 << 16) | (d7 << 8) | d8
                    return result

    def read_continuous(self):
        while True:
            if self.ser.read()[0] == 170:
                if self.ser.read()[0] == 170:
                    d1, d2, d3, d4, d5, d6, d7, d8 = list(self.ser.read(8))
                    result = (d5 << 24) | (d6 << 16) | (d7 << 8) | d8
                    yield result


class TempSensor(object):
    """
    A Temp Sensor Class.

    sensor = TempSensor()
    data = sensor.read()
    """

    def __init__(self):
        self.lib_main = CDLL("/home/pi/libabc.so")

    def read(self):
        return self.lib_main.getTemp()/100


class MPU6050(object):
    def read(self):
        return 10


class Example(QWidget):
    measure_finish_signal = pyqtSignal()
    temp_measure_finish_signal = pyqtSignal()
    motion_measure_finish_signal = pyqtSignal()

    def __init__(self, heart_max_length=2000, refresh_num=100, temp_mex_length=100, motion_max_length=1000):
        super(Example, self).__init__()
        self.initUI()
        self.generate_image()
        self.timer = QTimer(self)
        self.heart_sensor = HeartSensor(dev='/dev/ttyUSB0', baud_rate=115200)
        self.temp_sensor = TempSensor()
        self.motion_sensor = MPU6050()

        # 心脏相关参数
        self.status = True                      # True means not to measure
        self.refresh_num, self.process_num, self.process_count = refresh_num, 20, 0
        self.heart_max_length = heart_max_length
        # 最多记录多少个心脏读数
        self.max_data, self.min_data = 60000, -1000
        # 心脏读数允许的范围
        self.measure_finish_signal.connect(self.plotData)
        self.measure_thread = threading.Thread(target=self.measureing_thread)   # 定义线程
        self.measure_thread.start()  # 让线程开始工作
        self.satart_time = time.time()
        self.heart_history, self.time_history = [0] * self.heart_max_length, []
        self.x_axis_data = np.arange(0, self.heart_max_length) / 500

        # 温度采集相关参数
        self.temp_status = True                  # True means not to measure
        self.temp_max_length = temp_mex_length
        self.temp_history = []
        self.temp_interval = 1
        self.temp_measure_finish_signal.connect(self.plotTempData)
        self.temp_thread = threading.Thread(target=self.measureing_temp_thread)
        self.temp_thread.start()

        # 运动采集相关
        self.motion_status = True
        self.motion_max_length = motion_max_length
        self.motion_history = []
        self.motion_interval = 0.002
        self.motion_measure_finish_signal.connect(self.plotMotionData)
        self.motion_measure_thread = threading.Thread(target=self.measureing_moiton_thread)
        self.motion_measure_thread.start()

    def initUI(self):
        self.setGeometry(200, 200, 1000, 800)
        self.setWindowTitle("心电图")
        self.gridLayout = QGridLayout(self)
        self.frame = QFrame(self)
        self.frame.setFrameShape(QFrame.Panel)
        self.frame.setFrameShadow(QFrame.Plain)
        self.frame.setLineWidth(2)
        self.frame.setStyleSheet("background-color:rgb(0,255,255);")
        self.label = QLabel(self)
        self.label.setText("rs time: ")
        self.label.setAlignment(Qt.AlignCenter)

        self.button = QPushButton(self)
        self.button.setText("测量心电")
        self.button.clicked.connect(self.btnClick)
        self.heart_processbutton = QPushButton(self)
        self.heart_processbutton.setText("心电分析")
        self.heart_processbutton.clicked.connect(self.qrs_detect)

        self.temp_button = QPushButton(self)
        self.temp_button.setText("测量温度")
        self.temp_button.clicked.connect(self.tempClick)

        self.motion_button = QPushButton(self)
        self.motion_button.setText("运动测量")
        self.motion_button.clicked.connect(self.motionClick)

        self.gridLayout.addWidget(self.frame, 0, 0, 1, 2)
        self.gridLayout.addWidget(self.label, 1, 0, 1, 1)
        self.gridLayout.addWidget(self.button, 1, 1, 1, 1)
        self.gridLayout.addWidget(self.temp_button, 2, 0, 1, 1)
        self.gridLayout.addWidget(self.motion_button, 2, 1, 1, 1)
        self.gridLayout.addWidget(self.heart_processbutton, 3, 0, 1, 2)
        self.setLayout(self.gridLayout)

    def generate_image(self):
        verticalLayout = QVBoxLayout(self.frame)
        win = pg.GraphicsLayoutWidget(self.frame)
        win2 = pg.GraphicsLayoutWidget(self.frame)
        verticalLayout.addWidget(win)
        verticalLayout.addWidget(win2)

        p = win.addPlot(title="心电图")
        p.showGrid(x=False, y=False)
        p.setLabel(axis="left", text="幅值")
        p.setLabel(axis="bottom", text="t / 秒")
        p.setTitle("EKG")
        p.addLegend()
        self.curve_q = p.plot(pen=(0, 0, 0), symbolBrush=(19, 234, 201), symbolPen='w', symbol='s',
                              symbolSize=14, name="Q")
        self.curve_r = p.plot(pen=(0, 0, 0), symbolBrush=(0, 128, 0), symbolPen='w', symbol='+',
                              symbolSize=14, name="R")
        self.curve_s = p.plot(pen=(0, 0, 0), symbolBrush=(237, 177, 32), symbolPen='w', symbol='o',
                              symbolSize=14, name="S")
        self.curve1 = p.plot(pen="w", name="EKG")

        p = win2.addPlot(title="温度图")
        p.showGrid(x=True, y=True)
        p.setLabel(axis="left", text="Temp ℃")
        p.setLabel(axis="bottom", text="t")
        p.setTitle("Temp")
        p.addLegend()
        self.curve2 = p.plot(pen="w", name="Temp")

        p = win2.addPlot(title="运动传感器图")
        p.showGrid(x=True, y=True)
        p.setLabel(axis="left", text="a")
        p.setLabel(axis="bottom", text="t")
        p.setTitle("Motion")
        p.addLegend()
        self.curve3 = p.plot(pen="w", name="Motion")

    def plotData(self):
        self.heart_history = self.heart_history[-self.heart_max_length:]
        self.time_history = self.time_history[-self.heart_max_length:]
        q, q_y, r, r_y, s, s_y, status = self.qrs_detect(self.heart_history)
        self.curve_q.setData(q, q_y)
        self.curve_r.setData(r, r_y)
        self.curve_s.setData(s, s_y)
        self.curve1.setData(x=self.x_axis_data, y=self.heart_history)
        rs_time = np.mean([s[i]-r[i] for i in range(len(r))])
        # self.detect_qrs_online(self.heart_history[-self.refresh_num:])

    def plotTempData(self):
        self.temp_history = self.temp_history[-self.temp_max_length:]
        self.curve2.setData(self.temp_history)

    def plotMotionData(self):
        self.motion_history = self.motion_history[-self.motion_max_length:]
        self.curve3.setData(self.motion_history)

    def btnClick(self):
        if self.status:
            self.button.setText("停测心电")
            self.status = not self.status
        else:
            self.button.setText("测量心电")
            self.status = not self.status
            self.satart_time = time.time()

    def tempClick(self):
        if self.temp_status:
            self.temp_button.setText("停测温度")
            self.temp_status = not self.temp_status
        else:
            self.curve2.setData([])
            self.temp_button.setText("测量温度")
            self.temp_status = not self.temp_status

    def motionClick(self):
        if self.motion_status:
            self.motion_button.setText("停测运动")
            self.motion_status = not self.motion_status
        else:
            self.curve3.setData([])
            self.motion_button.setText("运动测量")
            self.motion_status = not self.motion_status

    def measureing_thread(self):
        while True:
            if self.status is True:
                time.sleep(0.08)
            else:
                temp_data, temp_time = [], []
                for data in self.heart_sensor.read_continuous():
                    if data < self.min_data or data > self.max_data:
                        continue
                    temp_data.append(data)
                    temp_time.append(time.time()-self.satart_time)
                    if len(temp_data) >= self.refresh_num:
                        self.heart_history.extend(temp_data)
                        self.time_history.extend(temp_time)
                        self.measure_finish_signal.emit()
                        temp_data, temp_time = [], []
                    if self.status:
                        break

    def measureing_temp_thread(self):
        while True:
            if self.temp_status:
                time.sleep(0.08)
            else:
                while True:
                    data = self.temp_sensor.read()
                    self.temp_history.append(data)
                    self.temp_measure_finish_signal.emit()
                    time.sleep(self.temp_interval)
                    if self.temp_status:
                        break

    def measureing_moiton_thread(self):
        while True:
            if self.motion_status:
                time.sleep(0.08)
            else:
                while True:
                    data = self.motion_sensor.read()
                    self.motion_history.append(data)
                    self.motion_measure_finish_signal.emit()
                    time.sleep(self.temp_interval)
                    if self.motion_status:
                        break

    def qrs_detect_raw(self):
        if len(self.heart_history) == self.heart_max_length:
            ecg = np.array(self.heart_history)
            time_ecg = np.array(self.time_history)
            fs = -int(self.heart_max_length/(time_ecg[self.heart_max_length-1]-time_ecg[0]))
            print(fs)
            R_peaks, S_pint, Q_point = EKG_QRS_detect(ecg, fs, QS=True, plot=False)
            r_y, s_y, q_y = ecg[R_peaks], ecg[S_pint], ecg[Q_point]
            self.curve_r.setData(x=R_peaks, y=r_y)
            self.curve_s.setData(x=S_pint, y=s_y)
            self.curve_q.setData(x=Q_point, y=q_y)
            heart_rate = int(s_y.shape[0]/(self.time_history[self.heart_max_length-1] - self.time_history[0])*60)
            a = "rs 时间: {:.4f} s ，心率：{:3d}".format(1/fs*np.mean(S_pint-R_peaks), heart_rate)
            self.label.setText(a)

    def qrs_detect(self, temp):
        stage = "Q"
        wave_yuzhi = 18000
        q_candi = 100000
        q_score = 0
        q_yuzhi = 10000
        r_candi = 0
        r_score = 0
        r_yuzhi = 10000
        s_candi = 100000
        s_score = 0
        s_yuzhi = 13

        q_list, r_list, s_list = [], [], []
        q_y_list, r_y_list, s_y_list = [], [], []

        for idx, item in enumerate(temp):
            if stage == "Q":
                if item < q_candi:
                    q_candi, q_candi_idx = item, idx
                    q_score = 0
                else:
                    q_score += (item - q_candi)
                if q_score > q_yuzhi:
                    r_candi, r_candi_idx = item, idx
                    q_score = 0
                    r_score = 0
                    stage = "R"
                    continue
            if stage == "R":
                if item > r_candi:
                    r_candi, r_candi_idx = item, idx
                else:
                    r_score += (r_candi - item)
                if r_score > r_yuzhi:
                    if r_candi - q_candi > wave_yuzhi:
                        s_candi, s_candi_idx = item, idx
                        r_score = 0
                        stage = "S"
                        continue
                    else:
                        stage = "Q"
                        continue
            if stage == "S":
                if item < s_candi:
                    s_candi, s_candi_idx = item, idx
                else:
                    s_score += 1
                if s_score > s_yuzhi:
                    s_score = 0
                    q_list.append(q_candi_idx/500)
                    r_list.append(r_candi_idx/500)
                    s_list.append(s_candi_idx/500)
                    print(q_candi_idx, r_candi_idx, s_candi_idx)
                    q_y_list.append(q_candi)
                    r_y_list.append(r_candi)
                    s_y_list.append(s_candi)
                    q_candi = 100000
                    r_candi = 0
                    s_candi = 100000
                    stage = "Q"
                    continue
        return q_list, q_y_list, r_list, r_y_list, s_list, s_y_list, stage


if __name__ == '__main__':
    app = QApplication(sys.argv)
    ex = Example()
    ex.show()
    sys.exit(app.exec_())
